from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import Optional, Literal
import uvicorn
import os
import tempfile
import requests
import base64
import torch
import numpy as np
from urllib.parse import urlparse
from torchaudio.transforms import Resample
import gushi_trainer
from run_gushi_eval import audioread_load # 复用命令行工具中的音频加载函数
import tts_service
import json
import qiniu

# --- FastAPI 应用设置 ---
app = FastAPI(
    title="中文古诗词评测API",
    description="基于FunASR的中文诗词朗读与背诵评测系统",
    version="1.0.0"
)

# --- 七牛云配置 ---
QINIU_ACCESS_KEY = os.getenv("QINIU_ACCESS_KEY")
QINIU_SECRET_KEY = os.getenv("QINIU_SECRET_KEY")
QINIU_BUCKET_DOMAIN = os.getenv("QINIU_BUCKET_DOMAIN")

# --- Pydantic模型定义 ---
class GushiEvalRequest(BaseModel):
    text: str
    task: Literal['reading', 'memorization'] = 'reading'
    audio_url: Optional[str] = None
    audio_base64: Optional[str] = None

    class Config:
        json_schema_extra = {
            "example": {
                "text": "床前明月光，疑是地上霜。",
                "task": "reading",
                "audio_url": "https://example.com/gushi.wav"
            }
        }

class TTSRequest(BaseModel):
    text: str
    voice: Optional[str] = None
    use_ssml: bool = True
    
    class Config:
        json_schema_extra = {
            "example": {
                "text": "床前明月光，疑是地上霜。举头望明月，低头思故乡。",
                "voice": "zhixiaoxia",
                "use_ssml": True
            }
        }

# --- 单例模式加载评测器 ---
# 在应用启动时加载模型，避免每次请求都重新加载，提高性能
@app.on_event("startup")
async def startup_event():
    print("服务器启动，开始加载评测模型...")
    app.state.trainer = gushi_trainer.get_gushi_trainer()
    print("评测模型加载完毕，服务准备就绪。")
    
    # 初始化TTS服务（需要配置阿里云密钥）
    # 注意：在生产环境中，请通过环境变量或配置文件提供这些密钥
    # app.state.tts_service = tts_service.get_tts_service(
    #     provider="aliyun",
    #     access_key_id=os.getenv("ALIYUN_ACCESS_KEY_ID"),
    #     access_key_secret=os.getenv("ALIYUN_ACCESS_KEY_SECRET"),
    #     app_key=os.getenv("ALIYUN_TTS_APP_KEY")
    # )

# --- 辅助函数 ---
def numpy_to_json_serializable(obj):
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {k: numpy_to_json_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [numpy_to_json_serializable(i) for i in obj]
    return obj

def process_audio(audio_path: str) -> torch.Tensor:
    signal_np, fs_orig = audioread_load(audio_path)
    if fs_orig == 0:
        raise HTTPException(status_code=400, detail="无法加载音频或音频为空")
    
    signal_torch = torch.from_numpy(signal_np.astype(np.float32))
    if fs_orig != 16000:
        resampler = Resample(orig_freq=fs_orig, new_freq=16000)
        signal_torch = resampler(signal_torch)
    
    return signal_torch.unsqueeze(0) if signal_torch.ndim == 1 else signal_torch

def download_audio(url: str) -> str:
    """
    下载音频文件。
    如果检测到是七牛云的URL且配置了AK/SK，则自动生成私有下载链接。
    """
    # 如果是七牛云地址且配置了密钥，则生成私有链接
    if QINIU_BUCKET_DOMAIN and url.startswith(QINIU_BUCKET_DOMAIN) and QINIU_ACCESS_KEY and QINIU_SECRET_KEY:
        try:
            auth = qiniu.Auth(access_key=QINIU_ACCESS_KEY, secret_key=QINIU_SECRET_KEY)
            signed_url = auth.private_download_url(url, expires=3600) # 1小时有效
            url = signed_url # 使用签名链接进行下载
            print(f"为七牛云链接生成签名URL: {url}")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"生成七牛云签名链接失败: {e}")

    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        suffix = os.path.splitext(urlparse(url).path)[1] or '.wav'
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
            tmp_file.write(response.content)
            return tmp_file.name
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=400, detail=f"下载音频文件失败: {e}")

def save_base64_audio(data: str) -> str:
    try:
        audio_data = base64.b64decode(data)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
            tmp_file.write(audio_data)
            return tmp_file.name
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"处理base64音频失败: {e}")

# --- API 端点 ---
@app.get("/health")
async def health_check():
    return {"status": "ok", "cuda_available": torch.cuda.is_available()}

@app.post("/eval_gushi")
async def eval_gushi_endpoint(request: GushiEvalRequest):
    if not request.audio_url and not request.audio_base64:
        raise HTTPException(status_code=400, detail="必须提供 audio_url 或 audio_base64")
    if request.audio_url and request.audio_base64:
        raise HTTPException(status_code=400, detail="不能同时提供 audio_url 和 audio_base64")

    tmp_path = None
    try:
        if request.audio_url:
            tmp_path = download_audio(request.audio_url)
        else:
            tmp_path = save_base64_audio(request.audio_base64)

        audio_tensor = process_audio(tmp_path)
        
        # 从应用状态中获取已加载的评测器
        trainer = app.state.trainer
        result = trainer.evaluate(audio_tensor, request.text, request.task)
        
        # 将结果中可能存在的numpy类型转换为JSON兼容类型
        serializable_result = numpy_to_json_serializable(result)
        
        return JSONResponse(content=serializable_result)

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {str(e)}")
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)

@app.post("/generate_demo")
async def generate_demo_reading(request: TTSRequest):
    """
    生成古诗词朗读示范音频
    """
    try:
        # 注意：这里需要您先配置TTS服务
        # 如果没有配置，返回提示信息
        if not hasattr(app.state, 'tts_service'):
            return JSONResponse(
                status_code=501,
                content={
                    "error": "TTS服务未配置",
                    "message": "请先配置阿里云TTS服务的API密钥",
                    "setup_guide": {
                        "step1": "获取阿里云AccessKey ID和Secret",
                        "step2": "申请智能语音交互服务的AppKey", 
                        "step3": "在环境变量中设置: ALIYUN_ACCESS_KEY_ID, ALIYUN_ACCESS_KEY_SECRET, ALIYUN_TTS_APP_KEY",
                        "step4": "重启服务即可使用TTS功能"
                    }
                }
            )
        
        # 生成朗读示范音频
        tts = app.state.tts_service
        audio_path = tts.generate_speech(
            text=request.text,
            voice=request.voice,
            use_ssml=request.use_ssml
        )
        
        # 将音频文件转换为base64返回
        with open(audio_path, 'rb') as f:
            audio_data = f.read()
            audio_base64 = base64.b64encode(audio_data).decode('utf-8')
        
        # 清理临时文件
        os.remove(audio_path)
        
        return JSONResponse(content={
            "status": "success",
            "audio_base64": audio_base64,
            "format": "wav",
            "text": request.text,
            "voice_used": request.voice or "zhixiaoxia"
        })
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"TTS生成失败: {str(e)}")

# --- 运行命令 ---
# uvicorn GUSHI.gushi_fastapi_app:app --host 0.0.0.0 --port 5002 --reload 